import torch
import torch.nn as nn
import torch.nn.functional as F


# https://github.com/kuangliu/pytorch-cifar/blob/master/models/lenet.py
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        out = F.relu(self.conv1(x))
        out = F.max_pool2d(out, 2)
        out = F.relu(self.conv2(out))
        out = F.max_pool2d(out, 2)
        out = out.view(out.size(0), -1)
        out = F.relu(self.fc1(out))
        out = F.relu(self.fc2(out))
        out = self.fc3(out)
        return out


class LeNet1(nn.Module):
    def __init__(self):
        super(LeNet1, self).__init__()
        in_dim = 3
        n_class = 10
        self.conv = nn.Sequential(nn.Conv2d(in_dim, 6, 5, stride=1, padding=0, ),
                                  nn.ReLU(),

                                  nn.ReLU(True),
                                  nn.MaxPool2d(2, 2),  # TODO padding
                                  nn.Conv2d(6, 16, 5, stride=1, padding=0),
                                  nn.ReLU(inplace=False),
                                  nn.MaxPool2d(2, 2))

        self.fc = nn.Sequential(
            nn.Linear(in_features=400, out_features=120, bias=True),
            nn.Linear(120, out_features=84),
            nn.Linear(84, n_class, False))
        out = torch.rand((10, 10, 10))
        out = out.view(out.size(0), -1)
        # print(out.shape)

    def forward(self, x):
        self.conv_nouse = nn.Conv2d(3, 6, 5)
        out = self.conv(x)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return out
